"""BLE Fragmentation Protocol for reliable data transfer.

Handles packet fragmentation, reassembly, CRC validation,
and retransmission for data larger than BLE MTU.
"""

import asyncio
import struct
import time
import logging
from enum import Enum, IntFlag
from typing import Dict, List, Optional, Callable, Any, Tuple
from dataclasses import dataclass, field
from collections import OrderedDict
import random

logger = logging.getLogger(__name__)


# Protocol constants
FRAME_MAGIC = 0xAA55
MAX_FRAME_SIZE = 20  # BLE MTU for BLE 4.0/4.1
HEADER_SIZE = 7  # Magic(2) + Seq(2) + Len(2) + Flags(1)
CRC_SIZE = 2
MAX_PAYLOAD_SIZE = MAX_FRAME_SIZE - HEADER_SIZE - CRC_SIZE  # 11 bytes


class FrameFlags(IntFlag):
    """Frame control flags."""
    FIRST_FRAME = 0x01     # First fragment
    LAST_FRAME = 0x02      # Last fragment
    ACK_REQUIRED = 0x04    # Requires acknowledgment
    RETRANSMIT = 0x08      # Retransmitted frame
    CANCEL = 0x10          # Cancel transfer
    ERROR = 0x20           # Error flag
    ACK = 0x40             # Acknowledgment frame
    NACK = 0x80            # Negative acknowledgment


class TransferState(Enum):
    """Transfer state machine states."""
    IDLE = 0
    SENDING = 1
    RECEIVING = 2
    WAITING_ACK = 3
    RETRANSMITTING = 4
    COMPLETE = 5
    ERROR = 6
    CANCELLED = 7


@dataclass
class Frame:
    """BLE data frame structure."""
    seq_num: int
    flags: FrameFlags
    payload: bytes
    length: Optional[int] = None
    crc: Optional[int] = None
    
    def __post_init__(self):
        """Calculate length and CRC if not provided."""
        if self.length is None:
            self.length = len(self.payload)
        if self.crc is None:
            self.crc = calculate_crc16(self.to_bytes(include_crc=False))
    
    def to_bytes(self, include_crc: bool = True) -> bytes:
        """Serialize frame to bytes."""
        data = struct.pack(
            '>HHHB',
            FRAME_MAGIC,
            self.seq_num,
            self.length,
            self.flags
        )
        data += self.payload
        
        # Pad payload if necessary
        if len(self.payload) < MAX_PAYLOAD_SIZE:
            data += b'\x00' * (MAX_PAYLOAD_SIZE - len(self.payload))
        
        if include_crc:
            data += struct.pack('>H', self.crc or 0)
        
        return data
    
    @classmethod
    def from_bytes(cls, data: bytes) -> 'Frame':
        """Deserialize frame from bytes."""
        if len(data) < MAX_FRAME_SIZE:
            raise ValueError(f"Invalid frame size: {len(data)}")
        
        magic, seq_num, length, flags = struct.unpack('>HHHB', data[:7])
        
        if magic != FRAME_MAGIC:
            raise ValueError(f"Invalid magic number: {magic:04X}")
        
        payload = data[7:7+length]
        crc = struct.unpack('>H', data[18:20])[0]
        
        # Verify CRC
        calculated_crc = calculate_crc16(data[:18])
        if crc != calculated_crc:
            raise ValueError(f"CRC mismatch: expected {calculated_crc:04X}, got {crc:04X}")
        
        return cls(
            seq_num=seq_num,
            flags=FrameFlags(flags),
            payload=payload,
            length=length,
            crc=crc
        )


@dataclass
class TransferStatistics:
    """Transfer performance statistics."""
    start_time: float = field(default_factory=time.time)
    end_time: Optional[float] = None
    bytes_sent: int = 0
    bytes_received: int = 0
    frames_sent: int = 0
    frames_received: int = 0
    retransmissions: int = 0
    errors: int = 0
    
    @property
    def duration(self) -> float:
        """Get transfer duration in seconds."""
        if self.end_time:
            return self.end_time - self.start_time
        return time.time() - self.start_time
    
    @property
    def throughput(self) -> float:
        """Get throughput in bytes per second."""
        duration = self.duration
        if duration > 0:
            return (self.bytes_sent + self.bytes_received) / duration
        return 0
    
    @property
    def error_rate(self) -> float:
        """Get error rate as percentage."""
        total_frames = self.frames_sent + self.frames_received
        if total_frames > 0:
            return (self.errors / total_frames) * 100
        return 0


def calculate_crc16(data: bytes) -> int:
    """Calculate CRC16-CCITT checksum.
    
    Args:
        data: Data to calculate checksum for
        
    Returns:
        16-bit CRC value
    """
    CRC16_POLY = 0x8005
    CRC16_INIT = 0xFFFF
    
    crc = CRC16_INIT
    for byte in data:
        crc ^= byte << 8
        for _ in range(8):
            if crc & 0x8000:
                crc = (crc << 1) ^ CRC16_POLY
            else:
                crc = crc << 1
            crc &= 0xFFFF
    
    return crc


class DataFragmenter:
    """Handles data fragmentation into BLE frames."""
    
    def __init__(self, ack_required: bool = True):
        """Initialize fragmenter.
        
        Args:
            ack_required: Whether frames require acknowledgment
        """
        self.ack_required = ack_required
        self.sequence_number = 0
    
    def fragment(self, data: bytes) -> List[Frame]:
        """Fragment data into frames.
        
        Args:
            data: Data to fragment
            
        Returns:
            List of frames
        """
        if not data:
            return []
        
        frames = []
        total_frames = (len(data) + MAX_PAYLOAD_SIZE - 1) // MAX_PAYLOAD_SIZE
        
        for i in range(total_frames):
            start = i * MAX_PAYLOAD_SIZE
            end = min(start + MAX_PAYLOAD_SIZE, len(data))
            payload = data[start:end]
            
            # Set frame flags
            flags = FrameFlags(0)
            if i == 0:
                flags |= FrameFlags.FIRST_FRAME
            if i == total_frames - 1:
                flags |= FrameFlags.LAST_FRAME
            if self.ack_required:
                flags |= FrameFlags.ACK_REQUIRED
            
            frame = Frame(
                seq_num=self.sequence_number,
                flags=flags,
                payload=payload
            )
            
            frames.append(frame)
            self.sequence_number = (self.sequence_number + 1) & 0xFFFF
        
        logger.debug(f"Fragmented {len(data)} bytes into {len(frames)} frames")
        return frames
    
    def reset_sequence(self):
        """Reset sequence number."""
        self.sequence_number = 0


class DataReassembler:
    """Handles frame reassembly into original data."""
    
    def __init__(self, timeout: float = 30.0):
        """Initialize reassembler.
        
        Args:
            timeout: Reassembly timeout in seconds
        """
        self.timeout = timeout
        self.buffers: Dict[int, Dict[int, Frame]] = {}  # Transfer ID -> frames
        self.expected_seq: Dict[int, int] = {}  # Transfer ID -> next expected seq
        self.start_times: Dict[int, float] = {}  # Transfer ID -> start time
        self.transfer_id = 0
    
    def add_frame(self, frame: Frame) -> Optional[bytes]:
        """Add frame to reassembly buffer.
        
        Args:
            frame: Frame to add
            
        Returns:
            Complete data if reassembly is complete, None otherwise
        """
        # Check for first frame to start new transfer
        if frame.flags & FrameFlags.FIRST_FRAME:
            self.transfer_id += 1
            self.buffers[self.transfer_id] = {}
            self.expected_seq[self.transfer_id] = frame.seq_num
            self.start_times[self.transfer_id] = time.time()
            logger.debug(f"Started new transfer {self.transfer_id}")
        
        # Find current transfer
        current_transfer = None
        for tid in self.buffers:
            if tid in self.expected_seq:
                # Check if this frame belongs to this transfer
                if frame.seq_num >= self.expected_seq[tid]:
                    current_transfer = tid
                    break
        
        if current_transfer is None:
            logger.warning(f"No active transfer for frame seq={frame.seq_num}")
            return None
        
        # Add frame to buffer
        self.buffers[current_transfer][frame.seq_num] = frame
        
        # Check if transfer is complete
        if frame.flags & FrameFlags.LAST_FRAME:
            return self._reassemble(current_transfer)
        
        return None
    
    def _reassemble(self, transfer_id: int) -> Optional[bytes]:
        """Reassemble frames for a transfer.
        
        Args:
            transfer_id: Transfer to reassemble
            
        Returns:
            Reassembled data or None if incomplete
        """
        if transfer_id not in self.buffers:
            return None
        
        frames = self.buffers[transfer_id]
        if not frames:
            return None
        
        # Sort frames by sequence number
        sorted_frames = OrderedDict(sorted(frames.items()))
        
        # Check for gaps
        expected = min(sorted_frames.keys())
        for seq_num in sorted_frames.keys():
            if seq_num != expected:
                logger.warning(f"Gap in sequence: expected {expected}, got {seq_num}")
                return None
            expected += 1
        
        # Reassemble data
        data = b''
        for frame in sorted_frames.values():
            data += frame.payload
        
        # Clean up
        del self.buffers[transfer_id]
        del self.expected_seq[transfer_id]
        del self.start_times[transfer_id]
        
        logger.debug(f"Reassembled transfer {transfer_id}: {len(data)} bytes")
        return data
    
    def cleanup_timeout(self):
        """Remove timed-out transfers."""
        current_time = time.time()
        expired = []
        
        for tid, start_time in self.start_times.items():
            if current_time - start_time > self.timeout:
                expired.append(tid)
        
        for tid in expired:
            logger.warning(f"Transfer {tid} timed out")
            if tid in self.buffers:
                del self.buffers[tid]
            if tid in self.expected_seq:
                del self.expected_seq[tid]
            if tid in self.start_times:
                del self.start_times[tid]


class RetransmissionManager:
    """Manages frame retransmission and acknowledgments."""
    
    def __init__(self, max_retries: int = 3, initial_timeout: float = 1.0):
        """Initialize retransmission manager.
        
        Args:
            max_retries: Maximum retransmission attempts
            initial_timeout: Initial timeout in seconds
        """
        self.max_retries = max_retries
        self.initial_timeout = initial_timeout
        self.pending_frames: Dict[int, Tuple[Frame, int, float]] = {}  # seq -> (frame, retries, timeout)
        self.ack_callbacks: Dict[int, Callable] = {}
        self.nack_callbacks: Dict[int, Callable] = {}
    
    async def send_with_retry(self, frame: Frame, 
                             send_func: Callable[[Frame], None]) -> bool:
        """Send frame with retransmission on timeout.
        
        Args:
            frame: Frame to send
            send_func: Function to send frame
            
        Returns:
            True if acknowledged, False if failed
        """
        retries = 0
        timeout = self.initial_timeout
        
        while retries <= self.max_retries:
            # Send frame
            if retries > 0:
                frame.flags |= FrameFlags.RETRANSMIT
                logger.debug(f"Retransmitting frame seq={frame.seq_num}, attempt {retries}")
            
            send_func(frame)
            
            # Wait for ACK
            ack_event = asyncio.Event()
            self.ack_callbacks[frame.seq_num] = lambda: ack_event.set()
            
            try:
                await asyncio.wait_for(ack_event.wait(), timeout)
                # ACK received
                del self.ack_callbacks[frame.seq_num]
                return True
                
            except asyncio.TimeoutError:
                # Timeout, retry
                retries += 1
                timeout = min(timeout * 2, 5.0)  # Exponential backoff with max 5s
        
        # Max retries exceeded
        logger.error(f"Failed to send frame seq={frame.seq_num} after {self.max_retries} retries")
        if frame.seq_num in self.ack_callbacks:
            del self.ack_callbacks[frame.seq_num]
        return False
    
    def handle_ack(self, seq_num: int):
        """Handle acknowledgment for a frame.
        
        Args:
            seq_num: Sequence number of acknowledged frame
        """
        if seq_num in self.ack_callbacks:
            logger.debug(f"ACK received for seq={seq_num}")
            self.ack_callbacks[seq_num]()
    
    def handle_nack(self, seq_num: int):
        """Handle negative acknowledgment for a frame.
        
        Args:
            seq_num: Sequence number of rejected frame
        """
        if seq_num in self.nack_callbacks:
            logger.debug(f"NACK received for seq={seq_num}")
            self.nack_callbacks[seq_num]()
    
    def create_ack_frame(self, seq_num: int) -> Frame:
        """Create ACK frame.
        
        Args:
            seq_num: Sequence number to acknowledge
            
        Returns:
            ACK frame
        """
        return Frame(
            seq_num=seq_num,
            flags=FrameFlags.ACK,
            payload=b''
        )
    
    def create_nack_frame(self, seq_num: int) -> Frame:
        """Create NACK frame.
        
        Args:
            seq_num: Sequence number to reject
            
        Returns:
            NACK frame
        """
        return Frame(
            seq_num=seq_num,
            flags=FrameFlags.NACK,
            payload=b''
        )


class TransferController:
    """Controls data transfer with progress tracking."""
    
    def __init__(self, fragmenter: DataFragmenter, 
                 reassembler: DataReassembler,
                 retransmission_mgr: RetransmissionManager):
        """Initialize transfer controller.
        
        Args:
            fragmenter: Data fragmenter
            reassembler: Data reassembler
            retransmission_mgr: Retransmission manager
        """
        self.fragmenter = fragmenter
        self.reassembler = reassembler
        self.retransmission_mgr = retransmission_mgr
        
        self.state = TransferState.IDLE
        self.statistics = TransferStatistics()
        self.progress_callback: Optional[Callable[[float], None]] = None
        self.completion_callback: Optional[Callable[[bool, bytes], None]] = None
        
        self.current_transfer_size = 0
        self.current_transfer_sent = 0
        self.cancel_event = asyncio.Event()
    
    async def send_data(self, data: bytes, 
                       send_func: Callable[[Frame], None]) -> bool:
        """Send data with fragmentation and retransmission.
        
        Args:
            data: Data to send
            send_func: Function to send individual frames
            
        Returns:
            True if successful, False otherwise
        """
        if self.state != TransferState.IDLE:
            logger.error("Transfer already in progress")
            return False
        
        self.state = TransferState.SENDING
        self.statistics = TransferStatistics()
        self.current_transfer_size = len(data)
        self.current_transfer_sent = 0
        self.cancel_event.clear()
        
        try:
            # Fragment data
            frames = self.fragmenter.fragment(data)
            total_frames = len(frames)
            
            # Send frames
            for i, frame in enumerate(frames):
                # Check for cancellation
                if self.cancel_event.is_set():
                    self.state = TransferState.CANCELLED
                    logger.info("Transfer cancelled")
                    return False
                
                # Send with retry
                self.state = TransferState.WAITING_ACK
                success = await self.retransmission_mgr.send_with_retry(frame, send_func)
                
                if not success:
                    self.state = TransferState.ERROR
                    self.statistics.errors += 1
                    logger.error(f"Failed to send frame {i+1}/{total_frames}")
                    return False
                
                # Update progress
                self.current_transfer_sent += len(frame.payload)
                self.statistics.frames_sent += 1
                self.statistics.bytes_sent += len(frame.payload)
                
                if self.progress_callback:
                    progress = (i + 1) / total_frames * 100
                    self.progress_callback(progress)
                
                self.state = TransferState.SENDING
            
            # Transfer complete
            self.state = TransferState.COMPLETE
            self.statistics.end_time = time.time()
            
            logger.info(f"Transfer complete: {len(data)} bytes in {total_frames} frames")
            logger.info(f"Statistics: {self.statistics.throughput:.2f} B/s, "
                       f"{self.statistics.error_rate:.2f}% errors")
            
            if self.completion_callback:
                self.completion_callback(True, data)
            
            return True
            
        except Exception as e:
            self.state = TransferState.ERROR
            self.statistics.errors += 1
            logger.error(f"Transfer error: {e}")
            
            if self.completion_callback:
                self.completion_callback(False, b'')
            
            return False
    
    def receive_frame(self, frame_data: bytes) -> Optional[bytes]:
        """Receive and process incoming frame.
        
        Args:
            frame_data: Raw frame data
            
        Returns:
            Complete data if transfer is complete, None otherwise
        """
        try:
            # Parse frame
            frame = Frame.from_bytes(frame_data)
            
            # Handle ACK/NACK
            if frame.flags & FrameFlags.ACK:
                self.retransmission_mgr.handle_ack(frame.seq_num)
                return None
            elif frame.flags & FrameFlags.NACK:
                self.retransmission_mgr.handle_nack(frame.seq_num)
                return None
            
            # Handle data frame
            if self.state == TransferState.IDLE:
                self.state = TransferState.RECEIVING
                self.statistics = TransferStatistics()
            
            self.statistics.frames_received += 1
            self.statistics.bytes_received += len(frame.payload)
            
            # Add to reassembler
            complete_data = self.reassembler.add_frame(frame)
            
            if complete_data:
                self.state = TransferState.COMPLETE
                self.statistics.end_time = time.time()
                
                logger.info(f"Received complete transfer: {len(complete_data)} bytes")
                
                if self.completion_callback:
                    self.completion_callback(True, complete_data)
                
                return complete_data
            
            return None
            
        except Exception as e:
            self.statistics.errors += 1
            logger.error(f"Error receiving frame: {e}")
            return None
    
    def cancel_transfer(self):
        """Cancel ongoing transfer."""
        if self.state in [TransferState.SENDING, TransferState.RECEIVING, 
                         TransferState.WAITING_ACK]:
            self.cancel_event.set()
            self.state = TransferState.CANCELLED
            logger.info("Transfer cancelled by user")
    
    def get_status(self) -> Dict[str, Any]:
        """Get transfer status.
        
        Returns:
            Status dictionary
        """
        progress = 0
        if self.current_transfer_size > 0:
            progress = (self.current_transfer_sent / self.current_transfer_size) * 100
        
        return {
            "state": self.state.name,
            "progress": int(progress),
            "transferred": self.statistics.bytes_sent + self.statistics.bytes_received,
            "throughput": self.statistics.throughput,
            "error_rate": self.statistics.error_rate,
            "retransmissions": self.statistics.retransmissions
        }


class FragmentationProtocol:
    """Main fragmentation protocol handler."""
    
    def __init__(self, ack_required: bool = True,
                 max_retries: int = 3,
                 reassembly_timeout: float = 30.0):
        """Initialize fragmentation protocol.
        
        Args:
            ack_required: Whether acknowledgments are required
            max_retries: Maximum retransmission attempts
            reassembly_timeout: Timeout for reassembly
        """
        self.fragmenter = DataFragmenter(ack_required)
        self.reassembler = DataReassembler(reassembly_timeout)
        self.retransmission_mgr = RetransmissionManager(max_retries)
        self.controller = TransferController(
            self.fragmenter,
            self.reassembler,
            self.retransmission_mgr
        )
        
        # Callbacks
        self.send_callback: Optional[Callable[[bytes], None]] = None
        self.receive_callback: Optional[Callable[[bytes], None]] = None
        
        logger.info("Fragmentation protocol initialized")
    
    async def send(self, data: bytes) -> bool:
        """Send data using fragmentation protocol.
        
        Args:
            data: Data to send
            
        Returns:
            True if successful
        """
        if not self.send_callback:
            logger.error("No send callback configured")
            return False
        
        def send_frame(frame: Frame):
            """Send frame using callback."""
            if self.send_callback:
                self.send_callback(frame.to_bytes())
        
        return await self.controller.send_data(data, send_frame)
    
    def receive(self, frame_data: bytes):
        """Receive frame data.
        
        Args:
            frame_data: Raw frame data
        """
        complete_data = self.controller.receive_frame(frame_data)
        
        if complete_data and self.receive_callback:
            self.receive_callback(complete_data)
    
    def set_callbacks(self, send_callback: Callable[[bytes], None],
                      receive_callback: Callable[[bytes], None]):
        """Set send and receive callbacks.
        
        Args:
            send_callback: Callback to send raw frame data
            receive_callback: Callback when complete data is received
        """
        self.send_callback = send_callback
        self.receive_callback = receive_callback
    
    def set_progress_callback(self, callback: Callable[[float], None]):
        """Set progress callback.
        
        Args:
            callback: Progress callback (0-100%)
        """
        self.controller.progress_callback = callback
    
    def cancel(self):
        """Cancel ongoing transfer."""
        self.controller.cancel_transfer()
    
    def get_status(self) -> Dict[str, Any]:
        """Get protocol status.
        
        Returns:
            Status dictionary
        """
        return self.controller.get_status()
    
    def reset(self):
        """Reset protocol state."""
        self.fragmenter.reset_sequence()
        self.controller.state = TransferState.IDLE
        logger.info("Protocol reset")


# Example usage
if __name__ == "__main__":
    import asyncio
    
    async def test_protocol():
        """Test fragmentation protocol."""
        # Create protocol instance
        protocol = FragmentationProtocol()
        
        # Simulated network with packet loss
        network_queue = asyncio.Queue()
        packet_loss_rate = 0.1  # 10% packet loss
        
        async def simulated_send(data: bytes):
            """Simulate sending over network."""
            if random.random() > packet_loss_rate:
                await network_queue.put(data)
                print(f"Sent frame: {len(data)} bytes")
            else:
                print(f"Packet lost! {len(data)} bytes")
        
        def receive_complete(data: bytes):
            """Handle complete data reception."""
            print(f"Received complete data: {len(data)} bytes")
            print(f"Data: {data[:50]}...")
        
        # Set callbacks
        protocol.send_callback = lambda d: asyncio.create_task(simulated_send(d))
        protocol.receive_callback = receive_complete
        
        # Test data
        test_data = b"Hello, BLE fragmentation!" * 10  # 260 bytes
        
        print(f"Sending {len(test_data)} bytes...")
        
        # Send data
        send_task = asyncio.create_task(protocol.send(test_data))
        
        # Simulate receiving frames
        async def receive_loop():
            while True:
                try:
                    frame_data = await asyncio.wait_for(network_queue.get(), 0.1)
                    protocol.receive(frame_data)
                except asyncio.TimeoutError:
                    continue
        
        receive_task = asyncio.create_task(receive_loop())
        
        # Wait for completion
        success = await send_task
        print(f"Send {'successful' if success else 'failed'}")
        
        # Get status
        status = protocol.get_status()
        print(f"Status: {status}")
        
        receive_task.cancel()
    
    # Run test
    asyncio.run(test_protocol())